use chrono::NaiveDateTime;
use futures::stream::StreamExt;
use sea_orm::{prelude::*, ActiveValue, QuerySelect, TransactionTrait};

use aspect_db_sql::{completed_job, pending_job};

use super::id::JobId;
use super::state;
use super::traits::Job;
pub use crate::errors::Error;

#[derive(Clone)]
pub struct JobDatastore {
    dbc: DatabaseConnection,
}

impl JobDatastore {
    pub fn new(dbc: DatabaseConnection) -> Result<Self, Error> {
        Ok(Self { dbc })
    }

    /// Low level function to create a job with an exact payload and queue name.
    async fn create_raw_job(
        &self,
        queue_name: String,
        created_at: NaiveDateTime,
        payload: Vec<u8>,
    ) -> Result<JobId, Error> {
        let id = Uuid::now_v7();

        let am = pending_job::ActiveModel {
            id: ActiveValue::Set(id),
            created_at: ActiveValue::Set(created_at),
            queue_name: ActiveValue::Set(queue_name),
            tries: ActiveValue::Set(0),
            last_try: ActiveValue::Set(None),
            next_try_after: ActiveValue::Set(None),
            payload: ActiveValue::Set(payload),
        };

        let res = am.insert(&self.dbc).await?;
        Ok(res.id)
    }

    /// Creates a new job for a certain job type, codec-encoding the payload.
    pub async fn create_job<J: Job>(
        &self,
        inst: &J,
        payload: &J::Payload,
    ) -> Result<state::JobMeta, Error> {
        let queue_name = inst.type_name().to_owned();
        let now = chrono::Utc::now().naive_utc();
        let payload_enc = aspect_codec::encode_to_vec(payload)?;
        let id = self
            .create_raw_job(queue_name.clone(), now, payload_enc)
            .await?;
        Ok(state::JobMeta::new(id, now, queue_name))
    }

    /// Moves a job entry from the pending jobs table into the completed table,
    /// optionally pruning the payload data in the process.
    pub async fn complete_job(
        &self,
        id: &Uuid,
        jcd: state::JobCompletionData,
        prune_payload: bool,
    ) -> Result<(), Error> {
        let tx = self.dbc.begin().await?;

        // First we have to query so that we know what to copy.
        let pj = pending_job::Entity::find_by_id(*id)
            .one(&tx)
            .await?
            .ok_or(Error::UnknownJob(*id))?;

        let completed_am = completed_job::ActiveModel {
            id: ActiveValue::Set(*id),
            created_at: ActiveValue::Set(pj.created_at),
            queue_name: ActiveValue::Set(pj.queue_name),
            completed_at: ActiveValue::Set(*jcd.completed_at()),
            run_duration: ActiveValue::NotSet,
            total_tries: ActiveValue::NotSet,
            status: ActiveValue::Set(jcd.status()),
            msg: ActiveValue::Set(jcd.into_msg()),
            payload: if prune_payload {
                ActiveValue::Set(Some(pj.payload))
            } else {
                ActiveValue::Set(None)
            },
        };

        // Now make both of the changes.
        completed_am.insert(&tx).await?;
        pending_job::Entity::delete_by_id(*id).exec(&tx).await?;

        // FIXME I suppose it's possible for us to do the copy entirely on the
        // database side in a single statement, but how do we express this in
        // terms of SeaORM?

        tx.commit().await?;
        Ok(())
    }

    /// Sets the job retry count and updates the `last_try` field to the current
    /// time.
    pub async fn update_job_retry_status(
        &self,
        id: &JobId,
        new_count: i16,
        retry_after: Option<NaiveDateTime>,
    ) -> Result<(), Error> {
        let now = chrono::Utc::now().naive_utc();

        let am = pending_job::ActiveModel {
            id: ActiveValue::Unchanged(*id),
            tries: ActiveValue::Set(new_count),
            last_try: ActiveValue::Set(Some(now)),
            next_try_after: ActiveValue::Set(retry_after),
            ..Default::default()
        };

        am.update(&self.dbc).await?;
        Ok(())
    }

    /// Queries jobs in a range for a specific queue / job type.
    pub async fn query_queue_jobs_in_range(
        &self,
        queue_name: &str,
        after: Option<&NaiveDateTime>,
        before: Option<&NaiveDateTime>,
    ) -> Result<Vec<state::PendingJobEntry>, Error> {
        let mut clause = pending_job::Column::QueueName.eq(queue_name);

        if let Some(after) = after {
            clause = clause.and(pending_job::Column::CreatedAt.gte(*after));
        }

        if let Some(before) = before {
            clause = clause.and(pending_job::Column::CreatedAt.lte(*before));
        }

        // Query as a stream to minimize the extra data we keep in memory, we
        // don't know how much this is going to be.
        let mut job_stream = pending_job::Entity::find()
            .filter(clause)
            .stream(&self.dbc)
            .await?;

        let mut entries = Vec::new();
        while let Some(j) = job_stream.next().await {
            let j = j?;
            let jm = state::JobMeta::new(j.id, j.created_at, j.queue_name);
            let jd = state::JobData::new(jm, j.payload);
            let pje = state::PendingJobEntry::new(jd, j.tries, j.last_try);
            entries.push(pje);
        }

        Ok(entries)
    }

    /// Queues ready jobs after some optional datetime, up to some limit.
    pub async fn query_ready_jobs(&self, limit: u64) -> Result<Vec<state::PendingJobEntry>, Error> {
        let now = chrono::Utc::now().naive_utc();

        // Same as above, use stream just in case.
        let mut job_stream = pending_job::Entity::find()
            .filter(
                pending_job::Column::NextTryAfter
                    .lte(now)
                    .or(pending_job::Column::NextTryAfter.is_null()),
            )
            .limit(limit)
            .stream(&self.dbc)
            .await?;

        let mut entries = Vec::new();
        while let Some(j) = job_stream.next().await {
            let j = j?;
            let jm = state::JobMeta::new(j.id, j.created_at, j.queue_name);
            let jd = state::JobData::new(jm, j.payload);
            let pje = state::PendingJobEntry::new(jd, j.tries, j.last_try);
            entries.push(pje);
        }

        Ok(entries)
    }

    /// Queries jobs that failed in the specified time range, optionally
    /// ignoring payloads.
    pub async fn query_failed_jobs_in_range(
        &self,
        after: Option<&NaiveDateTime>,
        before: Option<&NaiveDateTime>,
        ignore_payloads: bool,
    ) -> Result<Vec<state::FailEntry>, Error> {
        // TODO make this not even deliver the payloads

        let mut clause = completed_job::Column::Status.ne(0);

        if let Some(after) = after {
            clause = clause.and(completed_job::Column::CompletedAt.gte(*after));
        }

        if let Some(before) = before {
            clause = clause.and(completed_job::Column::CompletedAt.lte(*before));
        }

        let mut job_stream = completed_job::Entity::find()
            .filter(clause)
            .stream(&self.dbc)
            .await?;

        let mut entries = Vec::new();
        while let Some(j) = job_stream.next().await {
            let j = j?;
            let jm = state::JobMeta::new(j.id, j.created_at, j.queue_name);
            let payload = if ignore_payloads { None } else { j.payload };
            let jcd = state::JobCompletionData::new(j.completed_at, j.status, j.msg);
            let fje = state::FailEntry::new(jm, payload, jcd);
            entries.push(fje);
        }

        Ok(entries)
    }

    pub async fn query_complete_jobs(
        &self,
        queue_name: &str,
        before: Option<&NaiveDateTime>,
        after: Option<&NaiveDateTime>,
    ) -> Result<Vec<JobId>, Error> {
        let mut clause = completed_job::Column::QueueName.eq(queue_name);

        if let Some(after) = after {
            clause = clause.and(completed_job::Column::CompletedAt.gte(*after));
        }

        if let Some(before) = before {
            clause = clause.and(completed_job::Column::CompletedAt.lte(*before));
        }

        // TODO figure out how to make this only return the IDs
        let mut job_stream = completed_job::Entity::find()
            .filter(clause)
            .column(completed_job::Column::Id)
            .stream(&self.dbc)
            .await?;

        let mut ids = Vec::new();
        while let Some(j) = job_stream.next().await {
            let j = j?;
            ids.push(j.id);
        }

        Ok(ids)
    }
}
